#include <iostream>
#include <fstream>
#include <sstream>
#include <string>
#include <thread>
#include <boost/random.hpp>
#include <boost/random/uniform_int.hpp>
#include <boost/random/cauchy_distribution.hpp>
#include <boost/random/normal_distribution.hpp>
#include <boost/random/uniform_real.hpp>

using namespace std;

const int cityNum = 48; // number of cities
const int antNum = 20; // number of ants
const int samsara = 1000; // generation

double resultsOfSamsara[samsara] = {0}; // e.g. resultsOfSamsara[5000] means the min_length ever found at generation 5000

struct City {
    int number = 0;
    double x = 0;
    double y = 0;
};

struct Route {
    City cityPath[cityNum];
    double length = 0;
};

struct Ant {
    Route route; // the route built by this ant
    bool cityVisited[cityNum] = {false}; // value false means the i city is not visited
};

Ant ants[antNum]; // ant colony
Route bestRoute; // best route ever found
City cityList[cityNum]; // the original city list read from text file
City randomCityList[cityNum]; // a random shuffle from original city list
double distance_matrix[cityNum][cityNum] = {0}; // distance[i][j] means the distance between city i an city j
City cityVisited[cityNum]; // cityVisited[i] = 1 means city i visited
double pheromones[cityNum][cityNum] = {0}; // pheromone on every route between all cities
double minDistance; // shortest length ever found
int visitedCities[cityNum] = {0}; // greedy route


/* function declarations */
double random_probability(double, double);
int random_int(int, int);
void read_file(char *);
double distance(City a, City b);
void init_distance_matrix();
double path_len(Route route);
void print_city(City);
void print_city_list(City[], int);
void print_route(Route);

Route greed_route(City startCity);
City closestCity(City);

void savePheromones();

void knuth_durstenfeld_shuffle (City cities[], int size);

void init_ants(); // initialize all the ants start cities
void build_trip(Ant &ant, double alpha, double beta); // build route for an ant
bool tripFinished(Ant ant);
City choose_next_city(Ant ant, double alpha, double beta);
void add_to_Route(Ant &ant, City city);
int route_cities_num (Route route);
void pheromone_evaporation(double rho);
void simple_accumulate_pheromone(); // update the pheromone matrix
void add_elitist_ant_pheromone(int e, int ant_index);
void add_rank_based_pheromone(int omega);
bool cmp_ants(Ant a, Ant b);
int get_best_route(); // return 1 if get a better route

void save_probability(string file_name, double selectProbability[], int size, int nowCityNumber);

void saveResultsToTxt(double res[], basic_string<char> fileName);

/* function implementations */

// get a random double between min and max

const double lowPro = 0.001; // low limit of probability
double random_probability(double min, double max) {

    boost::mt19937 randomGenerator(time(0)*rand());
    boost::uniform_real<> uniformRealDistribution(min, max);
    boost::variate_generator<boost::mt19937, boost::uniform_real<>> variateGenerator(randomGenerator, uniformRealDistribution);
    return variateGenerator();
}

// get a random int between min and max
int random_int(int min, int max) {
    boost::mt19937 randomGenerator(time(0)*rand());
    boost::uniform_int<> uniformRealDistribution(min, max);
    boost::variate_generator<boost::mt19937, boost::uniform_int<>> variateGenerator(randomGenerator, uniformRealDistribution);
    return variateGenerator();
}

void read_file(char * dir_str) {

    FILE * fp;
    fp = fopen(dir_str, "r");
    if (fp == NULL) {
        cout<< "Open file " << dir_str << " failed!"<<endl;
        exit(1);
    }
    for (auto & city : cityList) {
        fscanf(fp, "%d", &city.number);
        fscanf(fp, "%lf", &city.x);
        fscanf(fp, "%lf", &city.y);
    }
    fclose(fp);

}

void print_city(City city) {
    cout<<"City No. "<<city.number<<" COORD: "<<city.x<<", "<<city.y<<endl;
}

void print_city_list(City list[], int len) {
    cout<<"city list: ";
    for (int i = 0; i < len; ++i) {
        if (list[i].number > 0)
            cout << list[i].number << ", ";
    }
    cout<<endl;
}

void print_route(Route route) {
    cout<<"Route: ";
    for (int i = 0; i < cityNum; ++i) {
        cout<<route.cityPath[i].number;
        if (i != cityNum-1) {
            cout<<"->";
        }

    }
    cout<<" Length: "<<route.length<<endl;
}

void save_route(string file_name, Route route) {
    stringstream ss;
    ss<<file_name;
    ofstream f_out(ss.str());
    if (!f_out.is_open())
    {
        cerr << "cannot create output file " << ss.str() << endl;
        exit(0);
    }

    for (auto & city : route.cityPath) {
        f_out << city.number <<"\t";
        f_out << cityList[city.number-1].x <<"\t";
        f_out << cityList[city.number-1].y <<"\t";
        f_out<<endl;
    }
    cout << ss.str() <<" created."<<endl;
    f_out.close();
}

// calculate the distance between a and b
double distance(City a, City b) {
    return ceil(sqrt((a.x-b.x)*(a.x-b.x)+(a.y-b.y)*(a.y-b.y)));
}

// calculate the length of the route
double path_len(Route route) {
    double path_length = 0;
    for (int i = 0; i < cityNum-1; ++i) {
        path_length += distance(route.cityPath[i], route.cityPath[i+1]);
    }
    path_length += distance(route.cityPath[0], route.cityPath[cityNum-1]);
    return path_length;
}

// calculate the distance between every 2 cities, and save them to a n*n matrix
void init_distance_matrix() {

    stringstream ss;
    ss<<"distance_matrix.txt";
    ofstream f_out(ss.str());

    if (!f_out.is_open())
    {
        cerr << "cannot create output file " << ss.str() << endl;
        exit(0);
    }

    for (int i = 0; i < cityNum; ++i) {
        for (int j = 0; j < cityNum; ++j) {
            distance_matrix[i][j] = distance(cityList[i], cityList[j]);
            f_out << distance_matrix[i][j] <<"\t";
        }
        f_out<<endl;
    }
    cout << ss.str() <<" created."<<endl;

    f_out.close();

}

// find a route by greedy algorithm
Route greed_route(City startCity) {
    Route route;
    for (int & visitedCity : visitedCities) {
        visitedCity = 0;
    }
    double sum_distance = 0;
    visitedCities[startCity.number-1] = 1; // add startCity to the route
    route.cityPath[0] = startCity;

    for (int i = 0; i < cityNum-1; ++i) {
        int closestIndex;
        double closestDistance = INT_MAX;
        for (int j = 0; j < cityNum; ++j) {
            if (visitedCities[j]==0 && (route.cityPath[i].number-1)!=j && closestDistance > distance_matrix[route.cityPath[i].number-1][j]) {
                closestDistance = distance_matrix[route.cityPath[i].number-1][j];
                closestIndex = j;
            }
        }
        route.cityPath[i+1] = cityList[closestIndex];
        visitedCities[closestIndex] = 1;
    }
    route.length = path_len(route);
    return route;
}

// save pheromones to a text
void savePheromones() {
    stringstream ss;
    ss<<"pheromones_matrix.txt";
    ofstream f_out(ss.str());
    if (!f_out.is_open())
    {
        cerr << "cannot create output file " << ss.str() << endl;
        exit(0);
    }

    for (int i = 0; i < cityNum; ++i) {
        for (int j = 0; j < cityNum; ++j) {
            f_out << pheromones[i][j] <<"\t";
        }
        f_out<<endl;
    }
    cout << ss.str() <<" created."<<endl;
    f_out.close();
}

// initialize the ant colony, put every ant at a random city as the start point
void init_ants() {
    for (auto & ant : ants) { // clear the ants colony info
        for (int m = 0; m < cityNum; ++m) {
            ant.route.cityPath[m].number = 0;
            ant.route.cityPath[m].x = 0;
            ant.route.cityPath[m].y = 0;
            ant.route.length = 0;
            ant.cityVisited[m] = false;
        }
    }

    knuth_durstenfeld_shuffle(randomCityList, cityNum);  // get a random sequence of cities in randomCityList

    if (antNum <= cityNum) { // when antNum <= cityNum
        for (auto & ant : ants) {
            ant.route.length = 0;
            for (int j = 0; j < cityNum; ++j) {
                ant.route.cityPath[j].number = 0;
                ant.cityVisited[j] = false;
            }
        }
        int i = 0;
        while (i < antNum) { // choose start city for each ant;
            ants[i].route.cityPath[0] = randomCityList[i];
            ants[i].cityVisited[ants[i].route.cityPath[0].number-1] = true;
            i++;
        }
    } else { // when antNum > cityNum
        cerr<<"init_ants() failed, cannot support condition when antNum > CityNum."<<endl;
        exit(-1);
    }
}

// build trip for the ant with parameters alpha and beta
void build_trip(Ant &ant, double alpha, double beta) {
    while (!tripFinished(ant)) { // choose next city to visit until trip finished
        City city = choose_next_city(ant, alpha, beta);
        add_to_Route(ant, city);
    }
    ant.route.length = path_len(ant.route);
//    print_route(ant.route);

}

// tell if the trip is finished, i.e. every city has already been visited
bool tripFinished(Ant ant) {
    bool flag = true;
    for (bool i : ant.cityVisited) {
        if (!i) {
            flag = i;
        }
    }
    return flag;
}

// choose a city from the unvisited city list as next city to go
City choose_next_city(Ant ant, double alpha, double beta) {

    int nowCityIndex = route_cities_num(ant.route)-1; // mark the number of cities has the ant traveled, notice that the index starts from 0
    if (nowCityIndex == -1) {
        cerr<<"The ant has not been set the start city.";
        exit(-1);
    }
    int nowCity = ant.route.cityPath[nowCityIndex].number;

    // roulette choose
    double denominator = 0;
    for (int i = 0; i < cityNum; ++i) {
        if (!ant.cityVisited[i]) {  // for unvisited cities
            denominator += pow(pheromones[nowCity-1][i], alpha) * pow(1.0/distance_matrix[nowCity-1][i], beta);
            if (distance_matrix[nowCity-1][i] == 0) {
                cout<<"nowCity-1: "<<nowCity-1<<" i: "<<i<<endl;
                exit(-1);
            }
        }
    }
    if (denominator == 0) {
        cerr<<"denominator is 0!"<<endl;
        denominator += 0.000000000001;
    }
    double numerator = 0;
    double sum = 0;
//    cout<<"selectProbability: "<<endl;
    double selectProbability[cityNum] = {0};
    for (int i = 0; i < cityNum; ++i) {
        if (!ant.cityVisited[i]) {  // for unvisited cities
            numerator = pow(pheromones[nowCity-1][i], alpha) * pow(1.0/distance_matrix[nowCity-1][i], beta);
            selectProbability[i] = numerator/denominator;
        }
//        cout<<i<<": "<<selectProbability[i]<<endl;
    }
    double accumulatedSelectProbability[ cityNum] = {0};
//    cout<<"accumulatedSelectProbability: "<<endl;
    for (int i = 0; i < cityNum; ++i) {
        for (int j = 0; j <= i; ++j) {
            if (!ant.cityVisited[i] && !ant.cityVisited[j]) {
                accumulatedSelectProbability[i] += selectProbability[j];
            }
        }
//        cout<<i+1<<": "<<accumulatedSelectProbability[i]<<endl;
    }
    double p;
    do {
        p = random_probability(0, 1);
    } while (p < lowPro); // make sure the p is not too small
//    cout<<"p: "<< p <<endl;
    int nextCity = -1;
    for (int i = 0; i < cityNum; ++i) {
        if (!ant.cityVisited[i]) {
            if (p < accumulatedSelectProbability[i]) {
                nextCity = i; // found the nextCity, break
                break;
            }
        }
    }
    if (nextCity == -1) {
        cerr<<"nextCity: "<<nextCity<<endl;
        exit(-1);
    }
    assert(nextCity != -1);
    return cityList[nextCity]; // return the city that has been found

}

void add_to_Route(Ant &ant, City city) {
    int cityIndex = route_cities_num(ant.route);
    ant.route.cityPath[cityIndex] = city;
    ant.cityVisited[city.number-1] = true;
}

int route_cities_num (Route route) {
    int count = 0;
    for (auto & city : route.cityPath) { // find the current city where this ant is at
        if (city.number != 0) {
            count++;
        }
    }
    return count;
}

void pheromone_evaporation(double rho) {
    for (int i = 0; i < cityNum; ++i) { // pheromones evaporation
        for (int j = 0; j < cityNum; ++j) {
            pheromones[i][j] = (1-rho) * pheromones[i][j];
        }
    }
}

// vanilla pheromone evaporation
void simple_accumulate_pheromone() {
    for (auto & ant : ants) { // pheromones accumulation
        double pheromoneToAdd = 1.0 / ant.route.length;
        for (int i = 0; i < cityNum; ++i) {
            if (i == cityNum-1) {
                pheromones[ant.route.cityPath[i].number-1][ant.route.cityPath[0].number-1] += pheromoneToAdd;
                pheromones[ant.route.cityPath[0].number-1][ant.route.cityPath[i].number-1] += pheromoneToAdd;
            } else {
                pheromones[ant.route.cityPath[i].number-1][ant.route.cityPath[i+1].number-1] += pheromoneToAdd;
                pheromones[ant.route.cityPath[i+1].number-1][ant.route.cityPath[i].number-1] += pheromoneToAdd;
            }
        }
    }
}

// elitist Ant
void elitistAntPheromone(int e, Ant ant) {
    double pheromoneToAdd = e * 1.0 / bestRoute.length;
    for (int i = 0; i < cityNum; ++i) {
        if (i == cityNum-1) {
            pheromones[ant.route.cityPath[i].number-1][ant.route.cityPath[0].number-1] += pheromoneToAdd;
            pheromones[ant.route.cityPath[0].number-1][ant.route.cityPath[i].number-1] += pheromoneToAdd;
        } else {
            pheromones[ant.route.cityPath[i].number-1][ant.route.cityPath[i+1].number-1] += pheromoneToAdd;
            pheromones[ant.route.cityPath[i+1].number-1][ant.route.cityPath[i].number-1] += pheromoneToAdd;
        }
    }
}// rank based
void rankBasedPheromone(int omega) {
    sort(ants, ants+antNum, cmp_ants);
    for (int k = 0; k < omega-1; ++k) {
        double pheromoneToAdd = (omega - k - 1) *  1.0 / ants[k].route.length;
        for (int i = 0; i < cityNum; ++i) {
            if (i == cityNum-1) {
                pheromones[ants[k].route.cityPath[i].number-1][ants[k].route.cityPath[0].number-1] += pheromoneToAdd;
                pheromones[ants[k].route.cityPath[0].number-1][ants[k].route.cityPath[i].number-1] += pheromoneToAdd;
            } else {
                pheromones[ants[k].route.cityPath[i].number-1][ants[k].route.cityPath[i+1].number-1] += pheromoneToAdd;
                pheromones[ants[k].route.cityPath[i+1].number-1][ants[k].route.cityPath[i].number-1] += pheromoneToAdd;
            }
        }
    }
}

// compare two ants by their route length
bool cmp_ants(Ant a, Ant b) {
    return a.route.length < b.route.length;
}

int get_best_route() { // when get a better route, return the index of ant, otherwise return -1
    double minLen = numeric_limits<double>::max();
    int flag = -1;
    for (int i = 0; i < antNum; ++i) {
        if (ants[i].route.length < minLen) {
            minLen = ants[i].route.length;
            flag = i;
        }
    }
    if (minDistance > minLen) {
        minDistance = minLen;
        bestRoute = ants[flag].route;
        return flag;
    }
    return flag;
}

void save_probability(string file_name, double selectProbability[], int size, int nowCityNumber) {
    stringstream ss;
    ss<<file_name;
    ofstream f_out(ss.str());
    if (!f_out.is_open())
    {
        cerr << "cannot create output file " << ss.str() << endl;
        exit(0);
    }
    f_out<<nowCityNumber<<endl;
    for (int i = 0; i < size; ++i) {
        f_out << selectProbability[i] << endl;
    }
    cout << ss.str() <<" created."<<endl;
    f_out.close();
}


/* knuth_durstenfeld_shuffle: gives me a random sequence of cities */
void knuth_durstenfeld_shuffle (City cities[], int size)
{
    for(int i = size-1; i >= 0; i--)
    {
        // guarantee the i'th value does not interfere with it's front value
        int index = random_int(0, i);
        swap(cities[index], cities[i]);
    }
}

void saveResultsToTxt(double res[], basic_string<char> fileName) {
    stringstream ss;
    ss<<fileName;
    ofstream f_out(ss.str());
    if (!f_out.is_open())
    {
        cerr << "cannot create output file " << ss.str() << endl;
        exit(0);
    }

    for (int i = 0; i < samsara; ++i) {
        f_out << res[i] <<"\t";
    }
    cout << ss.str() <<" created."<<endl;
    f_out.close();
}

int main() {
    clock_t start_t = clock();


    minDistance = numeric_limits<double>::max();

    read_file("att48.tsp.txt");
    init_distance_matrix();
    for (int i = 0; i < cityNum; ++i) {
        cityVisited[i].number = -1; // mark all city not visited
        randomCityList[i] = cityList[i]; // initialize the randomCityList as the cityList
    }

    Route greedRoute = greed_route(cityList[0]); // get a route by greedy algorithm
    double tau_0 = antNum/greedRoute.length;
    for (int i = 0; i < cityNum; ++i) { // initialize pheromones on all edges
        for (int j = 0; j < cityNum; ++j) {
            if (j == i) {
                pheromones[i][j] = 0;
            } else {
                pheromones[i][j] = tau_0;
            }
        }
    }
    savePheromones(); // save the initial pheromones

    double alpha = 0.99;
    double beta = 2.5;
    double rho = 0.5;


/* ==================== the ACO loop for one iteration ==================== */
    for (int l = 0; l < 1; ++l) {
        for (double & i : resultsOfSamsara) {
            i = 0;
        }
        for (int i = 0; i < samsara ; ++i) { // the samsara begins...
            init_ants();
            for (auto & ant : ants) {
                build_trip(ant, alpha, beta);
            }
            pheromone_evaporation(rho);
//            rankBasedPheromone(6);
            simple_accumulate_pheromone();
            int flag = get_best_route();
            if (flag != -1) {
//                elitistAntPheromone(cityNum, ants[flag]);
            }
            resultsOfSamsara[i] = minDistance;
            cout<<"samsara: "<< i <<" minDistance: "<<minDistance<<endl;

        }
//        string fileName = "minDis"+to_string(l) + ".txt";

        string fileName = "results\\minDisAS.txt";

        saveResultsToTxt(resultsOfSamsara, fileName);

        // reset data
        for (double & res : resultsOfSamsara) {
            res = 0;
        }
        minDistance = numeric_limits<double>::max();
        for (auto & c : bestRoute.cityPath) {
            c.number = 0;
            c.x = 0;
            c.y = 0;
            bestRoute.length = numeric_limits<double>::max();
        }
        for (int m = 0; m < cityNum; ++m) {
            for (int n = 0; n < cityNum; ++n) {
                pheromones[m][n] = tau_0;
            }
        }
        rho += 0.1;
    }
/* ==================== the ACO loop for one iteration ==================== */


    cout<<"<======================== procedure done. cost time: "<<double(clock()-start_t)/CLOCKS_PER_SEC<<"s ========================>"<<endl;
    return 0;
}

